import os
import numpy as np
import cv2 as cv
from PyCmpltrtok.common import *
from PyCmpltrtok.common_np import *
from python_ai.category.TensorRT.trt_my_helper_class import MyTrtHelper


if '__main__' == __name__:

    def _main():
        trt_path = '_models/fcn-resnet101.trt'

        sep('Load image')
        img_path = '/var/asuspei/my_src/TensorRT/TensorRT.git/quickstart/SemanticSegmentation/input.ppm'
        M = 1
        img = cv.imread(img_path, cv.IMREAD_COLOR)
        # cv.imshow('img', img)
        img_ = img.copy()
        img = img[:, :, ::-1]

        img = uint8_to_flt_by_lut(img)
        # check_np_detailed(img, 'img')
        # check_np_detailed(img[:, :, 0], 'img[:, :, 0]')
        # check_np_detailed(img[:, :, 1], 'img[:, :, 1]')
        # check_np_detailed(img[:, :, 2], 'img[:, :, 2]')
        # mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        # std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        # img -= mean
        # img /= std
        # check_np_detailed(img, 'img')
        # check_np_detailed(img[:, :, 0], 'img[:, :, 0]')
        # check_np_detailed(img[:, :, 1], 'img[:, :, 1]')
        # check_np_detailed(img[:, :, 2], 'img[:, :, 2]')
        img = img.transpose(2, 0, 1)
        input_arr = np.expand_dims(img, axis=0)
        check_np_detailed(input_arr, 'input_arr')

        sep('Load model by my helper class')
        print(trt_path)
        myObj = MyTrtHelper(
            trt_path, np.float32, input_arr.shape, [M, 1, *img_.shape[:2]],
            is_dynamic_shape=True,
            output_dtype=np.int32
        )

        sep('Run it')
        output_arr = myObj.predict(input_arr)
        check_np_detailed(output_arr, 'output_arr')

        # cv.waitKey()
        # cv.destroyAllWindows()

    _main()
